import glob
import pandas as pd
import datetime
import os
import re
import ast
from django.utils import timezone
from ai.models import Article, ArticleCategory, ArticleKeyword
from logger.logger_handler import Logger

logger = Logger(logger='import_data.py').get_logger()


def reset_articles():
    Article.objects.all().delete()
    ArticleKeyword.objects.all().delete()
    ArticleCategory.objects.all().delete()


def get_data_files():
    current_dir = os.getcwd()
    return glob.glob(os.path.join(current_dir, '*.xlsx'))


class ArticleImporter:
    def __init__(self):
        self.logger = logger
        # self.logger.setLevel(logging.NOTSET)
        self.utc_timezone = datetime.timezone.utc

    def find_or_create_keywords(self, keyword):
        res = ArticleKeyword.objects.filter(text__exact=keyword)
        if len(res):
            return res[0].id
        else:
            instance = ArticleKeyword.objects.create(text=keyword)
            return instance.id

    def find_or_create_category(self, category_name):
        res = ArticleCategory.objects.filter(name=category_name)

        if len(res):
            return res[0]
        else:
            instance = ArticleCategory.objects.create(name=category_name)
            return instance

    def convert_keyword_list(self, original_list):
        keywords = ast.literal_eval(original_list)
        return keywords

    def get_keywords_id(self, keywords_str):
        keywords = self.convert_keyword_list(keywords_str)
        foreign_keys = []
        for keyword in keywords:
            foreign_keys.append(self.find_or_create_keywords(keyword))
        return foreign_keys

    def parse_category(self, category_str, article):
        first_category = self.find_or_create_category(category_str)
        article.category.add(first_category)

    def import_data(self, file_name):
        # 避免读取到空字符串时读出的是 nan
        xlsx_df = pd.read_excel(file_name, keep_default_na=False)
        try:
            add_or_create_articles_keywords_set = set()

            for index in range(len(xlsx_df)):

                public_time = xlsx_df.loc[index, '时间']
                title = xlsx_df.loc[index, '标题']
                url = xlsx_df.loc[index, '标题链接']
                thumb = xlsx_df.loc[index, '缩略图']
                # datetime_obj = datetime.datetime.strptime(public_time, "%Y-%m-%d %H:%M:%S")
                public_naive_time = datetime.datetime.strptime(public_time, "%Y-%m-%d %H:%M:%S")
                public_time = timezone.make_aware(public_naive_time)
                content = xlsx_df.loc[index, '正文内容']
                abstract = xlsx_df.loc[index, '摘要']
                keywords_raw = xlsx_df.loc[index, '关键词']
                keyword_contents = xlsx_df.loc[index, '关键内容']
                categories_raw = xlsx_df.loc[index, '分类']
                risk = Article.NO_RISK if xlsx_df.loc[index, '风险判断'] != Article.HAS_RISK else Article.HAS_RISK
                author = xlsx_df.loc[index, '作者']
                article_source = xlsx_df.loc[index, '来源']
                article_instance = Article.objects.filter(url__exact=url).first()

                is_article_exists = True if article_instance else False
                if is_article_exists:
                    article = article_instance
                    article.title = title
                    article.url = url
                    article.thumb = thumb
                    article.public_time = public_time
                    article.content = content
                    article.abstract = abstract
                    article.keywords_raw = keywords_raw
                    article.keyword_contents = keyword_contents
                    article.categories_raw = categories_raw
                    # 风险的结果不能更新，因为飞书会回传设置操作结果，不能覆盖
                    # article.risk = risk
                    article.author = author
                    article.article_source = article_source
                else:
                    article = Article(
                        title=title,
                        url=url,
                        thumb=thumb,
                        public_time=public_time,
                        content=content,
                        abstract=abstract,
                        keywords_raw=keywords_raw,
                        keyword_contents=keyword_contents,
                        categories_raw=categories_raw,
                        risk=risk,
                        author=author,
                        article_source=article_source
                    )

                article.save()
                # 设置关键词
                article.keywords.set(self.get_keywords_id(keywords_raw))
                self.parse_category(categories_raw, article)
                article_keywords = article.keywords.all()
                add_or_create_articles_keywords_set.update(article_keywords)
            #  导入文章的同时需要同步统计文章关键词出现次数
            for keyword in add_or_create_articles_keywords_set:
                keyword.value = keyword.article_set.count()
                keyword.save()
            self.logger.info(f"导入文章: {len(xlsx_df.values)}篇")
        except Exception as e:
            self.logger.critical(f"发生异常：{str(e)}")
            raise e
